{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from k_means import KMeans\n",
    "from rv import RandomVariable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MulivariateGaussianMixture(RandomVariable):\n",
    "    \"\"\"\n",
    "    p(x|mu, L, pi(coef))\n",
    "    = sum_k pi_k N(x|mu_k, L_k^-1)\n",
    "    \"\"\"\n",
    "    def __init__(self,n_components,\n",
    "                nu=None,\n",
    "                cov=None,\n",
    "                tau=None,\n",
    "                coef=None,):\n",
    "        super().__init__()\n",
    "        assert isinstance(n_components,int)\n",
    "        self.n_components=n_components\n",
    "        self.mu=mu\n",
    "        if cov is not None and tau is not None:\n",
    "            raise ValueError(\"cannot assign both cov and tau at sametime\")\n",
    "        elif cov is not None:\n",
    "            self.cov=cov\n",
    "        elif tau is not None:\n",
    "            self.tau=tau\n",
    "        else:\n",
    "            self.cov=None\n",
    "            self.tau=None\n",
    "        self.coef=coef\n",
    "    \n",
    "    @property\n",
    "    def mu(self):\n",
    "        return self.parameter['mu']\n",
    "    \n",
    "    @mu.setter\n",
    "    def mu(self,mu):\n",
    "        if isinstance(mu,np.ndarray`):\n",
    "            assert mu.ndim==2\n",
    "            assert np.size(mu,0)==self.n_components#row numbers\n",
    "            self.ndim=np.size(mu,1)\n",
    "            self.parameter[\"mu\"]=mu\n",
    "        elif mu is None:\n",
    "            self.parameter[\"mu\"]=None\n",
    "        else:\n",
    "            raise TypeError(\"mu must be either np.ndarry or None\")\n",
    "            \n",
    "    @property\n",
    "    def cov(self):\n",
    "        return self.parameter[\"cov\"]\n",
    "    \n",
    "    @cov.setter\n",
    "    def cov(self,cov):\n",
    "        if isinstance(cov,np.ndarray):\n",
    "            assert cov.shape=(self.n_components,self.ndim,self.ndim)\n",
    "            self._tau=np.linalg.inv(cov)\n",
    "            self.parameter[\"cov\"]=cov\n",
    "        elif cov is None:\n",
    "            self.parameter[\"cov\"]=None\n",
    "            self._tau=None\n",
    "        else:\n",
    "            raise TypeError(\"cov must be either np.ndarray or None\")\n",
    "    \n",
    "    @property\n",
    "    def tau(self):\n",
    "        return self._tau\n",
    "    \n",
    "    @tau.setter\n",
    "    def tau(self,tau):\n",
    "        if isinstance(tau,np,ndarray):\n",
    "            assert tau.shape=(self.n_components,self.ndim,self.ndim)\n",
    "            self.cov=np.linalg.inv(tau)\n",
    "            self._tau=tau\n",
    "        elif tau is None:\n",
    "            self.parameter[\"cov\"]=None\n",
    "            self._tau=None\n",
    "        else:\n",
    "            raise TypeError(\"tau must be either np.ndarray or None\")\n",
    "    \n",
    "    @property\n",
    "    def coef(self):\n",
    "        return self.parameter[\"coef\"]\n",
    "    \n",
    "    @coef.setter\n",
    "    def coef(self,coef):\n",
    "        if isinstance(coef,np,ndarray):\n",
    "            assert coef.ndim==1\n",
    "            if np.isnan(coef).any():\n",
    "                self.parameter[\"coef\"]=np.ones(self.n_components)/self.n_components\n",
    "            elif not np.allclose(coef.sum(),1):\n",
    "                raise  ValueError(f\"sum of coef must be equal to 1 {coef}\")\n",
    "            self.parameter[\"coef\"]=coef\n",
    "        elif coef is None:\n",
    "            self.parameter[\"coef\"]=None\n",
    "        else:\n",
    "            raise TypeError(\"coef must be either np.ndarray or None\")\n",
    "    \n",
    "    @property\n",
    "    def shape(self):\n",
    "        if hasattr(self.mu,\"shape\"):\n",
    "            return self.mu.shape[1]\n",
    "        else:\n",
    "            return None\n",
    "    \n",
    "    def _gauss(self.X):\n",
    "        d=X[:,None,:]-self.mu\n",
    "        D_sq = np.sum(np.einsum('nki,kij->nkj', d, self.cov) * d, -1)\n",
    "        return(np.exp(-0.5*D_sq)\n",
    "               /np.sqrt(np.linalg.det(self.cov)*(2*np.pi)**self.ndim))\n",
    "    \n",
    "    def _fit(self,X):\n",
    "        cov=np.cov(X.T)\n",
    "        kmeans=KMeans(self.n_components)\n",
    "        kmeans.fit(X)\n",
    "        self.mu=kmeans.centers\n",
    "        self.cov=np.array([cov for _ in rang(self.n_components)])\n",
    "        self.coef=np.ones(self.n_components)/self.n_components\n",
    "        params=np.hstack(self.mu.ravel(),self.cov.ravel(),self.coef.ravel())\n",
    "        \n",
    "        while True:\n",
    "            stats=self._expectation(X)\n",
    "            self._maximization(X,stats)\n",
    "            new_params=np.hstack(self.mu.ravel(),self.cov.ravel(),self.coef.ravel())\n",
    "            if np.allclose(params,new_params):\n",
    "                break\n",
    "            else:\n",
    "                params=new_params\n",
    "    def _expectation(self,X):\n",
    "        gamma=self.coef%self._gauss(X)\n",
    "        gamma/=gamma.sum(axis=-1,keepdims=True)\n",
    "        return gamma\n",
    "    \n",
    "    def _maximization(self,X,gamma):\n",
    "        N_k=np.sum(gamma,axis=0)\n",
    "        self.coef=N_k/len(X)\n",
    "        self.mu=(X.T@gamma/N_k).T\n",
    "        d=x[:,None,:]-self.mu\n",
    "        self.cov=np.einsum('nki,nkj->kij', d, d * resps[:, :, None])/N_k[:,None,None]\n",
    "    \n",
    "    def joint_proba(self,X):\n",
    "        return self.coef*self._gauss(X)\n",
    "    \n",
    "    def _pdf(self,X):\n",
    "        return(np.sum(joint_proba(self,X),axis=-1))\n",
    "    \n",
    "    def classify_proba(self, X):\n",
    "        return self._expectation(X)\n",
    "    \n",
    "    def classify(self,X):\n",
    "        return np.argmax(self.classify_proba(X),axis=1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
